"""
代码尝试专用脚本
- 用于调试、测试新功能或片段
"""

import time
from tqdm import tqdm
import pandas as pd
import logging
from netcore import preprocess_city_df, fill_mode, build_network, draw_network
import yaml

# 读取 config.yaml
with open('config.yaml', 'r', encoding='utf-8') as f:
    config = yaml.safe_load(f)

CSV_PATH = config.get('csv_path')
FIELDS_CORE = config.get('score_fields')
CORE_NAME_MAP = config.get('score_fields_map', {})
CONTROL_VARS = config.get('control_vars')
FIGURES_DIR = config.get('figures_dir', 'figures')
THRESHOLD = config.get('threshold', 0.1)

try:
    df = pd.read_csv(CSV_PATH, low_memory=False)
    logging.info(f"CSV文件 {CSV_PATH} 读取成功，数据量：{df.shape[0]} 行，{df.shape[1]} 列。")
except Exception as e:
    logging.error(f"CSV文件读取失败：{e}")
    raise

if "城乡" not in df.columns:
    logging.error("未找到城乡字段，无法筛选城市样本。")
    raise ValueError("未找到城乡字段")

city_df = preprocess_city_df(df)
logging.info(f"城市样本筛选后行数: {len(city_df)}")

# 独热编码区县
if "区县" in city_df.columns:
    district_dummies = pd.get_dummies(city_df["区县"], prefix="区县")
    city_df = pd.concat([city_df, district_dummies], axis=1)
    control_vars = [v for v in CONTROL_VARS if v != "区县"] + list(district_dummies.columns)
else:
    control_vars = CONTROL_VARS.copy()

all_vars = list(set(FIELDS_CORE + control_vars))
df_net = city_df[FIELDS_CORE + control_vars].copy()
df_net = df_net.dropna(subset=FIELDS_CORE)
df_net = fill_mode(df_net, control_vars)

# 只绘制核心变量显著且相关系数绝对值大于阈值的网络图
data = df_net[FIELDS_CORE].dropna()
G = build_network(data, FIELDS_CORE, CORE_NAME_MAP, threshold=THRESHOLD)
draw_network(G, title="核心变量网络图", save_path=f"{FIGURES_DIR}/核心变量网络图.png", ch_font=None)

if __name__ == "__main__":
    print("="*40)
    start_time_str = time.strftime('%Y-%m-%d %H:%M:%S')
    start_time = time.time()
    print(f"⏱ 代码尝试开始 {start_time_str}")
    print("="*40)

    # 在此处编写你的测试代码
    for _ in tqdm(range(5), colour='yellow'):
        time.sleep(0.2)

    end_time_str = time.strftime('%Y-%m-%d %H:%M:%S')
    end_time = time.time()
    print("="*40)
    print(f"⌛️ 代码尝试结束 {end_time_str}")
    duration = end_time - start_time
    print(f"总耗时：{int(duration)} 秒（{duration/60:.2f} 分钟）")
    print("="*40)
